import numpy as np
from math import atan2, sqrt, cos, sin

class RobotArmModule:
    def __init__(self, l1=114, l2=60, mass=0.001):
        """
        初始化双连杆机械臂模块
        :param l1: 第一连杆长度
        :param l2: 第二连杆长度
        :param mass: 末端负载质量 (kg)
        """
        self.l1 = l1
        self.l2 = l2
        self.mass = mass
        self.g = 9.81  # 重力加速度 m/s²
        self.load_force = mass * self.g  # 负载力 (N)
        
    def normalize_angle(self, angle_deg):
        """
        将角度标准化到 -180° 到 180° 之间
        :param angle_deg: 输入角度 (度)
        :return: 标准化后的角度 (度)
        """
        normalized = angle_deg % 360  # 先转换到 0-360°
        if normalized > 180:
            normalized -= 360
        return normalized
    
    def degrees_to_radians(self, angle_deg):
        """将角度从度转换为弧度"""
        return np.radians(angle_deg)
    
    def radians_to_degrees(self, angle_rad):
        """将角度从弧度转换为度"""
        return np.degrees(angle_rad)
    
    def forward_kinematics(self, theta1_deg, theta2_deg):
        """
        正运动学计算：根据关节角度计算末端位置
        :param theta1_deg: 关节1角度 (度)
        :param theta2_deg: 关节2角度 (度)
        :return: 末端位置 (x, y)
        """
        # 标准化角度
        theta1_norm = self.normalize_angle(theta1_deg)
        theta2_norm = self.normalize_angle(theta2_deg)
        
        # 转换为弧度
        theta1_rad = self.degrees_to_radians(theta1_norm)
        theta2_rad = self.degrees_to_radians(theta2_norm)
        total_angle = theta1_rad + theta2_rad
        
        # 计算位置
        x1 = self.l1 * cos(theta1_rad)
        y1 = self.l1 * sin(theta1_rad)
        
        x2 = x1 + self.l2 * cos(total_angle)
        y2 = y1 + self.l2 * sin(total_angle)
        
        return (x2, y2)
    
    def inverse_kinematics(self, x, y, solution_type=0):
        """
        逆运动学计算：根据末端位置计算关节角度
        :param x: 目标x坐标
        :param y: 目标y坐标
        :param solution_type: 解类型 0:肘上解, 1:肘下解
        :return: (theta1_deg, theta2_deg) 或 None（无解时）
        """
        # 检查目标点是否在可达范围内
        d = sqrt(x**2 + y**2)
        max_reach = self.l1 + self.l2
        min_reach = abs(self.l1 - self.l2)
        
        if d > max_reach or d < min_reach:
            return None
        
        # 计算逆运动学
        cos_theta2 = (x**2 + y**2 - self.l1**2 - self.l2**2) / (2 * self.l1 * self.l2)
        cos_theta2 = max(min(cos_theta2, 1), -1)  # 处理数值计算误差
        
        # 两种可能的解
        theta2_rad_1 = -np.arccos(cos_theta2)  # 肘上
        theta2_rad_2 = np.arccos(cos_theta2)   # 肘下
        
        # 计算对应的theta1
        alpha = atan2(y, x)
        beta_1 = atan2(self.l2 * sin(theta2_rad_1), self.l1 + self.l2 * cos(theta2_rad_1))
        beta_2 = atan2(self.l2 * sin(theta2_rad_2), self.l1 + self.l2 * cos(theta2_rad_2))
        
        theta1_rad_1 = alpha - beta_1
        theta1_rad_2 = alpha - beta_2
        
        # 转换为度并标准化
        theta1_deg = [self.normalize_angle(self.radians_to_degrees(theta1_rad_1)),
                      self.normalize_angle(self.radians_to_degrees(theta1_rad_2))]
                      
        theta2_deg = [self.normalize_angle(self.radians_to_degrees(theta2_rad_1)),
                      self.normalize_angle(self.radians_to_degrees(theta2_rad_2))]
        
        return (theta1_deg[solution_type], theta2_deg[solution_type])
    
    def calculate_torques(self, theta1_deg, theta2_deg):
        """
        计算关节力矩
        :param theta1_deg: 关节1角度 (度)
        :param theta2_deg: 关节2角度 (度)
        :return: (torque1, torque2) 关节1和关节2的力矩 (N·m)
        """
        # 标准化角度并转换为弧度
        theta1_norm = self.normalize_angle(theta1_deg)
        theta2_norm = self.normalize_angle(theta2_deg)
        
        theta1_rad = self.degrees_to_radians(theta1_norm)
        theta2_rad = self.degrees_to_radians(theta2_norm)
        total_angle = theta1_rad + theta2_rad
        
        # 计算关节2的力矩 (N·m)
        torque2 = self.load_force * self.l2 * np.sin(np.pi/2 - total_angle)
        
        # 计算关节1的力矩 (N·m)
        torque1 = self.load_force * (
            self.l1 * np.sin(np.pi/2 - theta1_rad) + 
            self.l2 * np.sin(np.pi/2 - total_angle)
        )
        
        return (torque1, torque2)
    
    def set_load_mass(self, mass):
        """设置末端负载质量 (kg)"""
        self.mass = mass
        self.load_force = mass * self.g
